#ifndef UNICODE
#define UNICODE
#endif

#include <windows.h>
#include <winternl.h>
#include <tlhelp32.h>

#define STATUS_INFO_LENGTH_MISMATCH 0xc0000004

#define SystemHandleInformation 16
#define ObjectBasicInformation 0
#define ObjectNameInformation 1

typedef NTSTATUS (NTAPI *_NtQuerySystemInformation)(
    ULONG SystemInformationClass,
    PVOID SystemInformation,
    ULONG SystemInformationLength,
    PULONG ReturnLength
   );
typedef NTSTATUS (NTAPI *_NtDuplicateObject)(
    HANDLE SourceProcessHandle,
    HANDLE SourceHandle,
    HANDLE TargetProcessHandle,
    PHANDLE TargetHandle,
    ACCESS_MASK DesiredAccess,
    ULONG Attributes,
    ULONG Options
   );
typedef NTSTATUS (NTAPI *_NtQueryObject)(
    HANDLE ObjectHandle,
    ULONG ObjectInformationClass,
    PVOID ObjectInformation,
    ULONG ObjectInformationLength,
    PULONG ReturnLength
   );

typedef struct _SYSTEM_HANDLE
{
    ULONG ProcessId;
    BYTE ObjectTypeNumber;
    BYTE Flags;
    USHORT Handle;
    PVOID Object;
    ACCESS_MASK GrantedAccess;
} SYSTEM_HANDLE, *PSYSTEM_HANDLE;

typedef struct _SYSTEM_HANDLE_INFORMATION
{
    ULONG HandleCount;
    SYSTEM_HANDLE Handles[1];
} SYSTEM_HANDLE_INFORMATION, *PSYSTEM_HANDLE_INFORMATION;

typedef enum _POOL_TYPE
{
    NonPagedPool,
    PagedPool,
    NonPagedPoolMustSucceed,
    DontUseThisType,
    NonPagedPoolCacheAligned,
    PagedPoolCacheAligned,
    NonPagedPoolCacheAlignedMustS
} POOL_TYPE, *PPOOL_TYPE;

typedef struct _OBJECT_TYPE_INFORMATION
{
    UNICODE_STRING Name;
    ULONG TotalNumberOfObjects;
    ULONG TotalNumberOfHandles;
    ULONG TotalPagedPoolUsage;
    ULONG TotalNonPagedPoolUsage;
    ULONG TotalNamePoolUsage;
    ULONG TotalHandleTableUsage;
    ULONG HighWaterNumberOfObjects;
    ULONG HighWaterNumberOfHandles;
    ULONG HighWaterPagedPoolUsage;
    ULONG HighWaterNonPagedPoolUsage;
    ULONG HighWaterNamePoolUsage;
    ULONG HighWaterHandleTableUsage;
    ULONG InvalidAttributes;
    GENERIC_MAPPING GenericMapping;
    ULONG ValidAccess;
    BOOLEAN SecurityRequired;
    BOOLEAN MaintainHandleCount;
    USHORT MaintainTypeList;
    POOL_TYPE PoolType;
    ULONG PagedPoolUsage;
    ULONG NonPagedPoolUsage;
} OBJECT_TYPE_INFORMATION, *POBJECT_TYPE_INFORMATION;

PVOID GetLibraryProcAddress(PSTR LibraryName, PSTR ProcName)
{
    return GetProcAddress(GetModuleHandleA(LibraryName), ProcName);
}

int DisplayError(LPCWSTR msg)
{
	MessageBox(NULL, msg, L"SDTabletPC Fixer", MB_ICONERROR);
	return 1;
}

int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, int nShowCmd)
{
    _NtQuerySystemInformation NtQuerySystemInformation = GetLibraryProcAddress("ntdll.dll", "NtQuerySystemInformation");
    _NtDuplicateObject NtDuplicateObject = GetLibraryProcAddress("ntdll.dll", "NtDuplicateObject");
    _NtQueryObject NtQueryObject = GetLibraryProcAddress("ntdll.dll", "NtQueryObject");

    NTSTATUS status;
    PSYSTEM_HANDLE_INFORMATION handleInfo;
    ULONG handleInfoSize = 0x10000;
    ULONG pid = 0;
    HANDLE processHandle;
    ULONG i;
	HANDLE hProcessSnap;
	PROCESSENTRY32 pe32;
	PVOID objectNameInfo = NULL;
	ULONG objectNameInfoSize = 0x1000;
	ULONG tries = 0;
	LPWSTR caption = L"SDTabletPC Leaky Handle Fix";

	/* we don't need to be high priority */
	SetPriorityClass(GetCurrentProcess(), PROCESS_MODE_BACKGROUND_BEGIN);

	// try to find the SDTabletPC.exe process 5 times, wait 20 seconds between tries
	for(tries = 0; tries < 5; tries++)
	{
		// Take a snapshot of all processes in the system.
		hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
		if(hProcessSnap == INVALID_HANDLE_VALUE)
		{
			DisplayError(L"Could not take snapshot of processes");
			return 1;
		}

		// Set the size of the structure before using it.
		pe32.dwSize = sizeof(PROCESSENTRY32);

		// Retrieve information about the first process,
		// and exit if unsuccessful
		if(!Process32First(hProcessSnap, &pe32))
		{
			DisplayError(L"Could not query first process");
			CloseHandle(hProcessSnap);          // clean the snapshot object
			return(FALSE);
		}

		// Now walk the snapshot of processes, and
		// display information about each process in turn
		do
		{
			if(_wcsicmp(L"SDTabletPC.exe", pe32.szExeFile) == 0)
			{
				pid = pe32.th32ProcessID;
				break;
			}
		}
		while(Process32Next(hProcessSnap, &pe32));

		CloseHandle(hProcessSnap);

		/* if we found it, bail out of the loop, otherwise sleep and try again */
		if(pid != 0)
			break;
		else
			Sleep(20000);
	}

	/* tried 5 times and failed, dump out */
	if(pid == 0)
	{
		DisplayError(L"Couldn't find SDTabletPC.exe process.  Make sure you have the latest STMicro driver installed from Acer's support website (support.acer.com).");
		return 1;
	}

	/* open the SDTabletPC.exe process */
    if (!(processHandle = OpenProcess(PROCESS_DUP_HANDLE, FALSE, pid)))
    {
		DisplayError(L"Could not open SDTabletPC.exe PID.");
        return 1;
    }

	handleInfo = malloc(handleInfoSize);
	objectNameInfo = malloc(objectNameInfoSize);

	while(TRUE)
	{
		/* NtQuerySystemInformation won't give us the correct buffer size, so we guess by doubling the buffer size. */
		while ((status = NtQuerySystemInformation(SystemHandleInformation, handleInfo, handleInfoSize, NULL)) == STATUS_INFO_LENGTH_MISMATCH)
			handleInfo = (PSYSTEM_HANDLE_INFORMATION)realloc(handleInfo, handleInfoSize *= 2);

		/* NtQuerySystemInformation stopped giving us STATUS_INFO_LENGTH_MISMATCH. */
		if (!NT_SUCCESS(status))
		{
			DisplayError(L"NtQuerySystemInformation failed.");
			return 1;
		}

		/* enumerate the handles */
		for (i = 0; i < handleInfo->HandleCount; i++)
		{
			SYSTEM_HANDLE handle = handleInfo->Handles[i];
			HANDLE dupHandle = NULL;
			UNICODE_STRING objectName;
			ULONG returnLength;

			/* Check if this handle belongs to the PID the user specified or if it has the access mask below which will hang NtQueryObject */
			if (handle.ProcessId != pid || handle.GrantedAccess == 0x0012019f)
				continue;

			/* Duplicate the handle so we can query it. */
			if(!NT_SUCCESS(NtDuplicateObject(processHandle, (HANDLE)handle.Handle, GetCurrentProcess(), &dupHandle, 0, 0, 0)))
				continue;

			/* query the handle for its name information */
			if (!NT_SUCCESS(NtQueryObject(dupHandle, ObjectNameInformation, objectNameInfo, objectNameInfoSize, &returnLength)))
			{
				objectNameInfoSize = returnLength;

				/* Reallocate the buffer and try again. */
				objectNameInfo = realloc(objectNameInfo, returnLength);
				if (!NT_SUCCESS(NtQueryObject(dupHandle, ObjectNameInformation, objectNameInfo, returnLength, NULL)))
				{
					CloseHandle(dupHandle);
					continue;
				}
			}

			/* Cast our buffer into an UNICODE_STRING. */
			objectName = *(PUNICODE_STRING)objectNameInfo;

			/* Kill the original handle */
			if (objectName.Length && _wcsicmp(objectName.Buffer, L"\\Device\\STHall") == 0)
				DuplicateHandle(processHandle, (HANDLE)handle.Handle, NULL, NULL, 0, FALSE, DUPLICATE_CLOSE_SOURCE);

			/* Kill the duplicated handle */
			CloseHandle(dupHandle);
		}

		/* pause for a little while */
		Sleep(5000);
	}
}
